import requests
import hashlib
import zipfile
import os
import time
import concurrent.futures
from itertools import product

# --- 配置 ---
# 改为SYMBOLS列表, 可以添加任意数量的交易对
SYMBOLS = ["SUIUSDT"]
# TIME_INTERVALS = [
#     "1h", "5m"
# ]
TIME_INTERVALS = [
    "30m", "5m", "1h", "4h"
]
START_DATE_STR = "2024-01"
END_DATE_STR = "2025-09"  # 包含此月份

# 设置最大并发线程数
# 可以根据你的网络情况和CPU性能调整。对于网络I/O密集型任务，可以设置得比CPU核心数多。
MAX_WORKERS = 10

# 请确保以下路径存在, 或者程序将尝试创建它们
# 使用原始字符串 (raw strings) 来处理Windows路径中的反斜杠
BASE_SAVE_PATH_ROOT = r"F:\personal\binance_klines\\"


# --- 辅助函数 (与原版相同, 无需修改) ---

def download_file(url, save_path):
    """从URL下载文件到指定路径。"""
    try:
        # 确保目标文件夹存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        response = requests.get(url, stream=True, timeout=60)
        response.raise_for_status()

        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        # print(f"成功下载: {save_path}") # 在任务函数中打印更详细的信息
        return True
    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 404:
            # 404是正常情况 (例如某个月份没有数据), 不打印为错误
            pass
        else:
            print(f"[{os.path.basename(save_path)}] HTTP错误, 下载失败: {e}")
    except requests.exceptions.RequestException as e:
        print(f"[{os.path.basename(save_path)}] 下载错误: {e}")
    except Exception as e:
        print(f"[{os.path.basename(save_path)}] 下载时发生未知错误: {e}")
    return False


def verify_checksum(zip_file_path, checksum_file_path):
    """对照CHECKSUM文件验证ZIP文件的SHA256校验和。"""
    if not os.path.exists(zip_file_path) or not os.path.exists(checksum_file_path):
        return False
    try:
        with open(checksum_file_path, 'r') as f:
            expected_hash = f.read().strip().split()[0]

        sha256 = hashlib.sha256()
        with open(zip_file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(65536), b''):
                sha256.update(chunk)
        calculated_hash = sha256.hexdigest()

        return calculated_hash == expected_hash
    except Exception as e:
        print(f"[{os.path.basename(zip_file_path)}] 验证校验和时出错: {e}")
        return False


def unzip_file(zip_file_path, extract_to_dir):
    """将ZIP文件解压到指定目录。"""
    try:
        os.makedirs(extract_to_dir, exist_ok=True)
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to_dir)
        return True
    except zipfile.BadZipFile:
        print(f"错误: {os.path.basename(zip_file_path)} 是损坏的ZIP文件")
    except Exception as e:
        print(f"解压 {os.path.basename(zip_file_path)} 时出错: {e}")
    return False


# --- 核心任务处理函数 (新的) ---

def process_download_task(symbol, interval, year_str, month_str):
    """
    处理单个文件下载、验证和解压的完整流程。
    这个函数将被每个线程独立调用。
    """
    # 为每个任务生成一个唯一的日志前缀，方便在多线程输出中辨认
    log_prefix = f"[{symbol}-{interval}-{year_str}-{month_str}]"

    # --- 1. 动态生成路径和URL ---
    base_url = f"https://data.binance.vision/data/futures/um/monthly/klines/{symbol}/{interval}/"
    zip_filename = f"{symbol}-{interval}-{year_str}-{month_str}.zip"
    checksum_filename = f"{zip_filename}.CHECKSUM"

    zip_url = base_url + zip_filename
    checksum_url = base_url + checksum_filename

    # 每个symbol有自己的保存目录
    symbol_base_path = os.path.join(BASE_SAVE_PATH_ROOT, symbol)
    zip_save_dir = os.path.join(symbol_base_path, "zip")
    extract_base_dir = symbol_base_path

    final_zip_save_path = os.path.join(zip_save_dir, zip_filename)
    checksum_save_path = os.path.join(zip_save_dir, checksum_filename)
    extract_dir_for_interval = os.path.join(extract_base_dir, interval)

    # --- 2. 检查文件是否已处理 ---
    expected_csv_filename = f"{symbol}-{interval}-{year_str}-{month_str}.csv"
    expected_csv_path = os.path.join(extract_dir_for_interval, expected_csv_filename)


    # --- 3. 下载和验证逻辑 (与原版类似, 但增加了日志前缀) ---
    print(f"{log_prefix} 开始处理...")

    # 下载ZIP和CHECKSUM
    zip_downloaded = download_file(zip_url, final_zip_save_path)
    if not zip_downloaded:
        # 404 Not Found 是常见情况, 不需要报告失败
        if not os.path.exists(final_zip_save_path):
            # print(f"{log_prefix} 文件在服务器上不存在 (404)。")
            return f"{log_prefix} Skipped (404 Not Found)"
        else:  # 下载过程中发生其他错误
            return f"{log_prefix} Failed (ZIP download error)"

    checksum_downloaded = download_file(checksum_url, checksum_save_path)
    if not checksum_downloaded:
        print(f"{log_prefix} CHECKSUM下载失败，无法验证。")
        return f"{log_prefix} Failed (Checksum download error)"

    # 验证
    if verify_checksum(final_zip_save_path, checksum_save_path):
        # 解压
        if unzip_file(final_zip_save_path, extract_dir_for_interval):
            print(f"{log_prefix} 成功下载并解压。")
            status = f"{log_prefix} Success"
        else:
            print(f"{log_prefix} 校验通过但解压失败。")
            status = f"{log_prefix} Failed (Unzip error)"
    else:
        print(f"{log_prefix} 校验和验证失败！")
        # 如果校验失败，删除下载的zip文件以供下次重试
        try:
            os.remove(final_zip_save_path)
        except OSError as e:
            print(f"{log_prefix} 无法删除校验失败的ZIP文件: {e}")
        status = f"{log_prefix} Failed (Checksum mismatch)"

    # 清理CHECKSUM文件
    if os.path.exists(checksum_save_path):
        try:
            os.remove(checksum_save_path)
        except OSError as e:
            print(f"{log_prefix} 无法删除CHECKSUM文件: {e}")

    return status


# --- 主逻辑 (重构为任务分发器) ---
def main():
    start_time = time.time()

    # --- 1. 生成所有需要处理的日期 ---
    dates = []
    start_year, start_month = map(int, START_DATE_STR.split('-'))
    end_year, end_month = map(int, END_DATE_STR.split('-'))

    year, month = start_year, start_month
    while year < end_year or (year == end_year and month <= end_month):
        dates.append((str(year), f"{month:02d}"))
        month += 1
        if month > 12:
            month = 1
            year += 1

    # --- 2. 创建所有任务的组合 ---
    # 使用 itertools.product 生成 (symbol, interval, (year, month)) 的所有组合
    tasks_to_run = list(product(SYMBOLS, TIME_INTERVALS, dates))
    total_tasks = len(tasks_to_run)
    print(f"总共需要处理 {total_tasks} 个文件任务。")
    print(f"使用 {MAX_WORKERS} 个并发线程。")

    # --- 3. 使用ThreadPoolExecutor执行任务 ---
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        # 使用 executor.map 或 submit
        # submit更灵活，可以处理返回结果
        future_to_task = {
            executor.submit(process_download_task, symbol, interval, date[0], date[1]): (symbol, interval, date)
            for symbol, interval, date in tasks_to_run
        }

        completed_count = 0
        for future in concurrent.futures.as_completed(future_to_task):
            task_info = future_to_task[future]
            try:
                result = future.result()  # 获取任务函数的返回值
                # 你可以在这里处理结果, 例如统计成功/失败数量
                # print(result) # 打印每个任务的最终状态
            except Exception as exc:
                print(f"任务 {task_info} 生成了异常: {exc}")

            completed_count += 1
            # 打印进度
            print(f"进度: {completed_count}/{total_tasks} ({completed_count / total_tasks:.2%})", end='\r')

    end_time = time.time()
    print(f"\n\n所有任务已完成，总耗时: {end_time - start_time:.2f} 秒。")


if __name__ == "__main__":
    main()
